liger-kernels: add device context manager - #1029
Conversation
|
/kernel-bot build-and-stage |
|
No kernels provided. Use |
|
/kernel-bot build-and-stage liger-kernels |
|
Build request processed. Command: Dispatched (1):
|
|
/kernel-bot build-and-stage liger-kernels |
|
Build request processed. Command: Dispatched (1):
|
sayakpaul
left a comment
sayakpaul
left a comment
There was a problem hiding this comment.
I think it applies to all kernels that don't have this fix as you also pointed out in Slack. But I guess we can treat on a case by case basis?
|
Let me hunt them all down as best as I can, I think it's low effort for big impact tbh so no need to delay on case by case imo |
|
Added it to all triton launches I could find |
|
/kernel-bot merge-and-upload liger-kernels |
|
Build request processed. Command: Merge result: PR is already merged. Continuing with build/upload. Dispatched (1):
|
Fixes linkedin#1303. Some Triton launches can receive tensors on a non-current GPU, for example under `device_map="auto"`. Triton uses the active device for launch and compile state, so the launch needs to enter the input tensor's device context. This adds a small shared `device_context` helper and applies it directly around the RMSNorm, SwiGLU, and GeGLU launches. This follows the same pattern as huggingface/kernels-community#1029. Checks: - `make checkstyle` - `python -m py_compile src/liger_kernel/ops/utils.py src/liger_kernel/ops/rms_norm.py src/liger_kernel/ops/swiglu.py src/liger_kernel/ops/geglu.py` - `git diff --check` --------- Co-authored-by: Vaibhav Jindal <vaibhav.jndl@gmail.com>
2 participants
It's affecting models that use device map auto (under multi gpu setting) and
use_kernels=True, e.g. RMS Norm gets corrupted. But this goes for any triton kernel launch so it is propogated to all of themFor more details, see linkedin/Liger-Kernel#1303